
import matplotlib.pyplot as plt
import numpy as np


base_list = [".","A","T","G","C","N"]

def get_image(shape,path):
    data = np.ones(shape)
    with open(path,'r') as f:
        for id,line in enumerate(f):
            if id %2 == 0:
                line = line.strip().split("\t")
                base,pos,quality_scores  = line[0],int(line[1]),[int(q)-1 for q in line[2:]]
            else:
                count = [int(c)-1 for c in line.strip().split("\t")[2:]]
                data[base_list.index(base),pos,quality_scores] = count
    return data

sample_info = [
    {
        'name':"build in",
        'shape':(6,53,40),
        'path':"Emp50R1.txt"
    },
    {
        'name':"sequencer",
        'shape':(6,36,40),
        'path':"GA2.txt"
    }
]

fig, ax = plt.subplots(5,2,sharex=True, sharey=True)
for id_smaple,info in enumerate(sample_info):
    data = get_image(info['shape'],info['path'])
    for id,base in enumerate(base_list[:-1]):
        Z = data[id]
        im = ax[id,id_smaple].imshow(Z[:36].T,cmap='magma',origin='lower',aspect='auto')
        ax[id,id_smaple].xlabel = 'test'
        ax[id,id_smaple].xaxis.set_major_locator(plt.MaxNLocator(8))
        ax[id,id_smaple].yaxis.set_major_locator(plt.MaxNLocator(3))
        
fig.colorbar(im, ax=ax, label='count')
plt.show()
